import torch
import torch.nn as nn
import numpy as np
import model
import pickle
import utils
import matplotlib.pyplot as plt

# Hyper Parameters
hyper_params = {
        "device": "cpu",
        "dropout": 0.0,
        "data_path": "data"
        }

print('predict: Loading data...')
# RSSI List
rssi_list: utils.RSSIList = pickle.load(open('./net/normal_rssi_list.pkl', 'rb'))
hyper_params["mac_cnt"] = rssi_list.mac_cnt
hyper_params["device_cnt"] = rssi_list.device_cnt
hyper_params["max_x"] = rssi_list.max_x
hyper_params["max_y"] = rssi_list.max_y

# Model
net = model.Model(hyper_params)
net.load_state_dict(torch.load('./net/net.pkl'))
print('predict: Finished!')

# ID Temp Number
id_rssi = {}
alpha = 0.8


# Preprocess Data
def preprocess(rssi_map: dict, name_id: str):
    input_tensor = torch.zeros([hyper_params['device_cnt']], dtype=torch.float)
    input_cnt = torch.zeros([hyper_params['device_cnt']], dtype=torch.float)

    for (k, v) in rssi_map.items():
        if rssi_list.mac_map.get(k) is None:
            print(f'Unknown MAC: {k}')
            continue
        input_tensor[rssi_list.mac_map[k]
                     ] += (v - rssi_list.rssi_min_value) / 70.0
        input_cnt[rssi_list.mac_map[k]] += 1.0

    if id_rssi.get(name_id) is not None:
        last_tensor, last_cnt = id_rssi[name_id]
        input_tensor += last_tensor
        input_cnt += last_cnt
    
    if torch.sum(input_cnt == 0) == 0:
        input_tensor /= input_cnt
        id_rssi[name_id] = (input_tensor, input_cnt)
        return 0
    else:
        id_rssi[name_id] = (input_tensor, input_cnt)
        return 1


# Predict Function
def predict(name_id: str, save_fig = True):
    input_tensor = id_rssi[name_id][0]

    with torch.no_grad():
        _, z = net.dae(input_tensor.unsqueeze(0))
        pred = torch.softmax(net.predict(z)[0], dim=-1).view(rssi_list.max_x, rssi_list.max_y).detach().cpu().numpy()
        
        # Fig
        if save_fig:
            plt.imshow(pred.copy(), cmap=plt.cm.hot)
            plt.colorbar()
            plt.title(name_id)
            plt.savefig('result.png')
            plt.close()

    # Clear Input Tensor
    id_rssi.pop(name_id)

    return pred
